package ai.zixing.mashibing.advance_class.first.class03;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * bfprt 算法
 * arr 上找第 k 小的数
 */
public class Code01_FindMinKth {

    public static class MaxHeapComparator implements Comparator<Integer> {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    }

    // 利用大根堆，时间复杂度O(N*logK)
    public static int minKth1(int[] arr, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(new MaxHeapComparator());
        // 先放入 k 个数
        for (int i = 0; i < k; i++) {
            maxHeap.add(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            if (arr[i] < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.add(arr[i]);
            }
        }
        return maxHeap.peek();
    }

    // 改写快排，时间复杂度O(N)
    public static int minKth2(int[] array, int k) {
        int[] arr = copyArray(array);
        return process2(arr, 0, arr.length - 1, k - 1);
    }

    public static int[] copyArray(int[] arr) {
        int[] ans = new int[arr.length];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = arr[i];
        }
        return ans;
    }


    // arr[L..R] 范围上，如果排序的话（不是真的去排序），找位于 index 的数
    public static int process2(int[] arr, int L, int R, int index) {
        if (L == R) {
            return arr[L];
        }
        // arr 不止有一个数

        // 等概率在L..R 上随机一个数
        int pivot = arr[L + (int) (Math.random() * (R - L + 1))];
        // 荷兰国旗问题
        int[] range = partition(arr, L, R, pivot);
        // arr[L..R] 的值都是 pivot
        if (index >= range[0] && index <= range[1]) {
            return arr[index];
            // 小于区去找
        } else if (index < range[0]) {
            return process2(arr, L, range[0] - 1, index);
            // 大于区去找
        } else {
            return process2(arr, range[1] + 1, R, index);
        }
    }

    // 荷兰国旗问题
    public static int[] partition(int[] arr, int L, int R, int pivot) {
        // 小于区结束下标记为 L - 1
        int less = L - 1;
        // 大于区的开始下标记为 R + 1
        int more = R + 1;
        int cur = L;
        while (cur < more) {
            if (arr[cur] < pivot) {
                // 当前数小于 pivot，当前数与小于区右边的数交换，当前数++
                swap(arr, ++less, cur++);
            } else if (arr[cur] > pivot) {
                // cur 与 大于区的左边的交换，当前位置不变
                swap(arr, cur, --more);
                // arr[cur] == pivot 遍历下一个
            } else {
                cur++;
            }
        }
        return new int[]{less + 1, more - 1};
    }

    public static void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    // arr[L..R] 如果排序的话，位于 index 位置的数是什么，返回

    public static int bfprt(int[] arr, int L, int R, int index) {
        if (L == R) {
            return arr[L];
        }
        int pivot = medianOfMedians(arr, L, R);
        int[] range = partition(arr, L, R, pivot);
        if (index >= range[0] && index <= range[1]) {
            return arr[index];
        } else if (index < range[0]) {
            return bfprt(arr, L, range[0] - 1, index);
        } else {
            return bfprt(arr, range[1] + 1, R, index);
        }
    }

    // arr[L..R] 五个数一组
    // 每个小组内部排序
    // 每个小猪中位数组成 marr
    // marr 中的中位数返回
    public static int medianOfMedians(int[] arr, int L, int R) {
        int size = R - L + 1;
        // 补齐组数
        int offset = size % 5 == 0 ? 0 : 1;
        int[] mArr = new int[size / 5 + offset];
        for (int team = 0; team < mArr.length; team++) {
            // 分组并计算每组中位数
            int teamFirst = L + team * 5;
            mArr[team] = getMedian(arr, teamFirst, Math.min(R, teamFirst + 4));
        }
        // marr 中找到中位数
        return bfprt(mArr, 0, mArr.length - 1, mArr.length / 2);
    }

    // 获取中位数
    public static int getMedian(int[] arr, int L, int R) {
        insertSort(arr, L, R);
        return arr[(L + R) / 2];
    }

    public static void insertSort(int[] arr, int L, int R) {
        for (int i = L + 1; i <= R; i++) {
            for (int j = i - 1; j >= L && arr[j] > arr[j + 1]; j--) {
                swap(arr, j, j + 1);
            }
        }
    }

    // 利用bfprt算法，时间复杂度O(N)
    public static int minKth3(int[] array, int k) {
        int[] arr = copyArray(array);
        return bfprt(arr, 0, arr.length - 1, k - 1);
    }

    public static int[] generateRandomArray(int maxSize, int maxValue) {
        int[] arr = new int[(int) (Math.random() * maxSize) + 1];
        for (int i = 0; i < arr.length; i++) {
            arr[i] = (int) (Math.random() * maxValue);
        }
        return arr;
    }

    public static void main(String[] args) {
        int testTime = 1000000;
        int maxSize = 100;
        int maxValue = 100;
        System.out.println("test begin");
        for (int i = 0; i < testTime; i++) {
            int[] arr = generateRandomArray(maxSize, maxValue);
            int k = (int) (Math.random() * arr.length) + 1;
            int ans1 = minKth1(arr, k);
            int ans2 = minKth2(arr, k);
            int ans3 = minKth3(arr, k);
            if (ans1 != ans2 || ans2 != ans3) {
                System.out.println("Oops");
            }
        }
        System.out.println("test finish");
    }
}
